/*
 * Copyright (c) 2010-2011 Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */
package jsr335.agent;

import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
import static org.objectweb.asm.Opcodes.ALOAD;
import static org.objectweb.asm.Opcodes.ATHROW;
import static org.objectweb.asm.Opcodes.DUP;
import static org.objectweb.asm.Opcodes.ILOAD;
import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static org.objectweb.asm.Opcodes.INVOKESTATIC;
import static org.objectweb.asm.Opcodes.IRETURN;
import static org.objectweb.asm.Opcodes.NEW;

import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;
import java.util.BitSet;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import org.objectweb.asm.AnnotationVisitor;
import org.objectweb.asm.Attribute;
import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Type;
import org.objectweb.asm.commons.EmptyVisitor;

/**
 * @author Remi Forax
 */
public class Agent {
    private static final String WOVEN_ANNOTATION = "Ljsr335/agent/Woven;";

    public static void premain(@SuppressWarnings("unused") String agentArgs,
                               Instrumentation instrumentation) {

        final HashMap<ClassKey, ClassMirror> classMap = new HashMap<ClassKey, ClassMirror>();

        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public synchronized byte[] transform(final ClassLoader classLoader,
                                                 final String className, Class<?> classBeingRedefined,
                                                 ProtectionDomain protectionDomain, byte[] classfileBuffer) {

                //System.out.println("transform " + className);

                try {
                    ClassReader classReader = new ClassReader(classfileBuffer);

                    ClassMirror classMirror = getClassMirror(classLoader, className,
                                                             classReader);

                    if (Modifier.isInterface(classReader.getAccess())) {
                        return null;
                    }

                    final Map<String, MethodMirror> internalMethodMap = classMirror
                            .getInternalMethodMap();

                    ClassWriter classWriter = new ClassWriter(classReader, 0);
                    classReader.accept(new ClassAdapter(classWriter) {
                        @Override
                        public void visitEnd() {
                            super.visitEnd();

                            Collection<MethodMirror> methods = internalMethodMap.values();
                            for (MethodMirror method : methods) {
                                DefenderAttribute attribute = method.getAttribute();
                                if (attribute == null) {
                                    continue;
                                }

                                String methodDesc = method.getDescriptor();
                                MethodVisitor mv = super.visitMethod(ACC_PUBLIC,
                                                                     method.getName(), methodDesc, method.getSignature(),
                                                                     method.getExceptions());

                                mv.visitCode();

                                if (attribute == DefenderAttribute.PROTOTYPE) {
                                    generateConflictCode(mv, methodDesc, method.getName() + method.getDescriptor());
                                    continue;
                                }

                                mv.visitVarInsn(ALOAD, 0);

                                int max = 1; // this = 1
                                Type[] types = Type.getArgumentTypes(methodDesc);
                                for (int i = 0; i < types.length; i++) {
                                    Type type = types[i];
                                    mv.visitVarInsn(type.getOpcode(ILOAD), i + 1);
                                    max += type.getSize();
                                }

                                mv.visitMethodInsn(INVOKESTATIC,
                                                   attribute.getClassName(), attribute.getName(),
                                                   attribute.getDescriptor());
                                mv.visitInsn(Type.getReturnType(methodDesc).getOpcode(IRETURN));

                                mv.visitMaxs(max, max);

                                AnnotationVisitor av = mv.visitAnnotation(WOVEN_ANNOTATION, true);
                                av.visit("definingInterface", "unknown");
                                av.visit("fromClass", attribute.getClassName());
                                av.visit("methodName", attribute.getName());
                                av.visit("methodSignature", attribute.getDescriptor());
                                av.visitEnd();

                                mv.visitEnd();
                            }
                        }

                        private void generateConflictCode(MethodVisitor mv, String methodDesc, String nameAndDesc) {
                            mv.visitTypeInsn(NEW, "java/lang/LinkageError");
                            mv.visitInsn(DUP);
                            mv.visitLdcInsn("extension method default conflict " + nameAndDesc);
                            mv.visitMethodInsn(INVOKESPECIAL, "java/lang/LinkageError", "<init>", "(Ljava/lang/String;)V");
                            mv.visitInsn(ATHROW);

                            mv.visitMaxs(3, Type.getArgumentsAndReturnSizes(methodDesc) >> 2);
                            mv.visitEnd();
                        }
                    }, 0);

                    byte[] byteArray = classWriter.toByteArray();
                    // CheckClassAdapter.verify(new ClassReader(byteArray), classLoader, false, new PrintWriter(System.err));

                    return byteArray;

                }
                catch (Throwable t) {
                    t.printStackTrace();
                    return null;
                }
            }

            private ClassMirror getClassMirror(ClassLoader classLoader,
                                               String className) {
                String fileName = className.replace('.', '/') + ".class";
                InputStream input = (classLoader == null) ?
                                    ClassLoader.getSystemResourceAsStream(fileName) :
                                    classLoader.getResourceAsStream(fileName);
                try {
                    return getClassMirror(classLoader, className, new ClassReader(input));
                }
                catch (IOException e) {
                    throw (NoClassDefFoundError) new NoClassDefFoundError().initCause(e);
                }
            }

            private ClassMirror getClassMirror(final ClassLoader classLoader,
                                               String className, ClassReader classReader) {
                ClassKey classKey = new ClassKey(classLoader, className);
                ClassMirror classMirror = classMap.get(classKey);
                if (classMirror != null) {
                    return classMirror;
                }

                final ClassMirror clazzMirror = new ClassMirror();

                final HashMap<String, MethodMirror> methodMap =
                        interfaceMap(clazzMirror, classLoader, classReader.getInterfaces());

                String superName = classReader.getSuperName();
                if (superName != null) {
                    ClassMirror superclassMirror = getClassMirror(classLoader, superName);
                    clazzMirror.addSupertype(superclassMirror);
                    methodMap.putAll(superclassMirror.getExportedMethodMap());
                }

                classReader.accept(new EmptyVisitor() {
                    @Override
                    public MethodVisitor visitMethod(int access, final String name,
                                                     final String desc, final String signature,
                                                     final String[] exceptions) {
                        final String nameAndDesc = name + desc;

                        final MethodMirror method = new MethodMirror(access,
                                                                     name, desc, signature, exceptions);
                        methodMap.put(nameAndDesc, method);

                        return new EmptyVisitor() {
                            @Override
                            public void visitAttribute(Attribute attr) {
                                DefenderAttribute attribute = (DefenderAttribute) attr;
                                method.setAttribute(attribute);
                            }
                        };
                    }
                }, new Attribute[]{DefenderAttribute.PROTOTYPE},
                                   ClassReader.SKIP_CODE);

                boolean isInterface = Modifier.isInterface(classReader.getAccess());
                clazzMirror.importMethods(isInterface, methodMap);
                classMap.put(classKey, clazzMirror);
                return clazzMirror;
            }

            HashMap<String, MethodMirror> interfaceMap(ClassMirror clazz, ClassLoader classLoader, String[] interfaces) {
                HashMap<String, MethodMirror> hashMap = new HashMap<String, MethodMirror>();

                int length = interfaces.length;
                ClassMirror[] mirrors = new ClassMirror[length];
                for (int i = 0; i < length; i++) {
                    mirrors[i] = getClassMirror(classLoader, interfaces[i]);
                }

                // remove super-interfaces
                BitSet marks = new BitSet(length);
                for (int i = 0; i < length; i++) {
                    if (marks.get(i))
                        continue;

                    ClassMirror cm = mirrors[i];

                    for (int j = 0; j < length; j++) {
                        if (i == j || marks.get(j))
                            continue;

                        if (cm.isSubtypeOf(mirrors[j])) {
                            marks.set(j);  // mark as should be skipped
                        }
                    }
                }

                for (int i = 0; i < length; i++) {
                    if (marks.get(i))
                        continue;

                    ClassMirror supertype = mirrors[i];
                    clazz.addSupertype(supertype);

                    Map<String, MethodMirror> methodMap = supertype.getExportedMethodMap();
                    for (Map.Entry<String, MethodMirror> entry : methodMap.entrySet()) {
                        MethodMirror method = entry.getValue();

                        String nameAndDesc = entry.getKey();
                        MethodMirror oldMethod = hashMap.get(nameAndDesc);
                        if (oldMethod == null || oldMethod.getAttribute() == null || method.getAttribute() == DefenderAttribute.PROTOTYPE) {
                            hashMap.put(nameAndDesc, method);
                            continue;
                        }

                        if (method.getAttribute() == null || oldMethod.getAttribute() == DefenderAttribute.PROTOTYPE) {
                            continue;
                        }

                        if (method.getAttribute().equals(oldMethod.getAttribute())) {
                            continue;
                        }

                        // mark as conflict
                        MethodMirror newMethod = oldMethod.erase();
                        newMethod.setAttribute(DefenderAttribute.PROTOTYPE);
                        hashMap.put(nameAndDesc, newMethod);
                    }
                }
                return hashMap;
            }

        }, true);
    }
}
